#include "vsf.h"

static vsf_err_t vsfip_tcp_forward_thread(struct vsfsm_pt_t *pt, vsfsm_evt_t evt)
{
	struct vsfip_forward_work_t *work = (struct vsfip_forward_work_t *)pt->user_data;
	struct vsfip_socket_t *rx_so = work->rx_so, *tx_so = work->tx_so;
	struct vsfsm_pt_t *caller_pt = &work->caller_pt;
	vsf_err_t err;

	vsfsm_pt_begin(pt);

	while (1)
	{
		caller_pt->state = 0;
		vsfsm_pt_entry(pt);
		err = vsfip_tcp_recv(caller_pt, evt, rx_so, &work->curbuf);
		if (err > 0) return VSFERR_NOT_READY; else if (err < 0) break;

		vsfip_netif_release(work->curbuf->netif);
		work->curbuf->netif = NULL;

		caller_pt->state = 0;
		vsfsm_pt_entry(pt);
		err = vsfip_tcp_send(caller_pt, evt, tx_so, work->curbuf, false);
		if (err > 0) return VSFERR_NOT_READY; else if (err < 0) break;
	}

	if (!*work->exit)
	{
		*work->exit = 1;

		caller_pt->state = 0;
		vsfsm_pt_entry(pt);
		err = vsfip_tcp_close(caller_pt, evt, rx_so);
		if (err > 0) return VSFERR_NOT_READY; else vsfip_close(rx_so);

		caller_pt->state = 0;
		vsfsm_pt_entry(pt);
		err = vsfip_tcp_close(caller_pt, evt, tx_so);
		if (err > 0) return VSFERR_NOT_READY; else vsfip_close(tx_so);

		if (vsfsm_sem_pend(work->sem, &work->sm))
			vsfsm_pt_wfe(pt, VSFSM_EVT_USER);
		vsf_bufmgr_free(work);

		vsfdbg_prints("tcp_forward: close connection" VSFCFG_DEBUG_LINEEND);
	}
	else
	{
		vsfsm_sem_post(work->sem);
	}

	vsfsm_pt_end(pt);
	return VSFERR_NONE;
}

static vsf_err_t vsfip_tcp_forward_listener_thread(struct vsfsm_pt_t *pt, vsfsm_evt_t evt)
{
	struct vsfip_tcp_forward_t *forwarder = (struct vsfip_tcp_forward_t *)pt->user_data;
	struct vsfsm_pt_t *caller_pt = &forwarder->listener.caller_pt;
	struct vsfip_socket_t *listener_so = forwarder->listener.so;
	struct vsfip_forward_conn_t *conn = forwarder->cur_conn;
	vsf_err_t err;

	vsfsm_pt_begin(pt);

again:
	while (!conn)
	{
		conn = (struct vsfip_forward_conn_t *)vsf_bufmgr_malloc(sizeof(*conn));
		if (!conn)
			vsfsm_pt_delay(pt, 100);
		else
		{
			memset(conn, 0, sizeof(*conn));
			forwarder->cur_conn = conn;
			if (!conn->fso)
			{
				conn->fso = vsfip_socket(AF_INET, IPPROTO_TCP);
				if (!conn->fso) vsfsm_pt_delay(pt, 100);
			}
		}
	}

	vsfdbg_prints("tcp_forward: tcp wait connection" VSFCFG_DEBUG_LINEEND);

	caller_pt->state = 0;
	vsfsm_pt_entry(pt);
	err = vsfip_tcp_accept(caller_pt, evt, listener_so, &conn->so);
	if (err > 0) return VSFERR_NOT_READY;
	else if (err < 0)
	{
		vsfdbg_printf("tcp_forward: tcp accept failed on port %d"
			VSFCFG_DEBUG_LINEEND, conn->so->remote_sockaddr.sin_port);
		goto fail;
	}

	vsfdbg_printf("tcp_forward: tcp acceptted on port %d"
			VSFCFG_DEBUG_LINEEND, conn->so->remote_sockaddr.sin_port);

	caller_pt->state = 0;
	vsfsm_pt_entry(pt);
	err = vsfip_tcp_connect(caller_pt, evt, conn->fso, &forwarder->sockaddr);
	if (err > 0) return VSFERR_NOT_READY;
	else if (err < 0)
	{
		vsfdbg_printf("tcp_forward: tcp connect failed on port %d"
			VSFCFG_DEBUG_LINEEND, forwarder->sockaddr.sin_port);
		goto fail;
	}

	vsfdbg_printf("tcp_forward: new connection on port %d"
			VSFCFG_DEBUG_LINEEND, forwarder->sockaddr.sin_port);
	forwarder->cur_conn = NULL;
	conn->list.next = forwarder->conn.next;
	forwarder->conn.next = &conn->list;

	vsfsm_sem_init(&conn->sem, 0, VSFSM_EVT_USER);

	conn->tx.rx_so = conn->so;
	conn->tx.tx_so = conn->fso;
	conn->tx.exit = &conn->exit;
	conn->tx.sem = &conn->sem;
	conn->tx.pt.thread = vsfip_tcp_forward_thread;
	conn->tx.pt.user_data = &conn->tx;
	conn->tx.caller_pt.sm = &conn->tx.sm;
	vsfsm_pt_init(&conn->tx.sm, &conn->tx.pt);

	conn->rx.rx_so = conn->fso;
	conn->rx.tx_so = conn->so;
	conn->rx.exit = &conn->exit;
	conn->rx.sem = &conn->sem;
	conn->rx.pt.thread = vsfip_tcp_forward_thread;
	conn->rx.pt.user_data = &conn->rx;
	conn->rx.caller_pt.sm = &conn->rx.sm;
	vsfsm_pt_init(&conn->rx.sm, &conn->rx.pt);

	conn = NULL;
	goto again;

fail:
	if (conn != NULL)
	{
		if (conn->so != NULL)
		{
			caller_pt->state = 0;
			vsfsm_pt_entry(pt);
			err = vsfip_tcp_close(caller_pt, evt, conn->so);
			if (err > 0) return VSFERR_NOT_READY; else if (err < 0) break;
			vsfip_close(conn->so);
		}
		if (conn->fso != NULL)
		{
			caller_pt->state = 0;
			vsfsm_pt_entry(pt);
			err = vsfip_tcp_close(caller_pt, evt, conn->fso);
			if (err > 0) return VSFERR_NOT_READY; else if (err < 0) break;
			vsfip_close(conn->fso);
		}
		vsf_bufmgr_free(conn);
	}

	vsfsm_pt_end(pt);
	return VSFERR_NONE;
}

vsf_err_t vsfip_tcp_forward_start(struct vsfip_tcp_forward_t *forwarder)
{
	struct vsfip_socket_t *so = vsfip_socket(AF_INET, IPPROTO_TCP);
	if (!so) return VSFERR_NOT_ENOUGH_RESOURCES;

	if (vsfip_bind(so, forwarder->sockaddr.sin_port) || vsfip_listen(so, 255))
		goto close_and_exit;

	forwarder->listener.so = so;
	forwarder->listener.pt.thread = vsfip_tcp_forward_listener_thread;
	forwarder->listener.pt.user_data = forwarder;
	forwarder->listener.caller_pt.sm = &forwarder->listener.sm;
	return vsfsm_pt_init(&forwarder->listener.sm, &forwarder->listener.pt);

close_and_exit:
	vsfip_close(so);
	return VSFERR_FAIL;
}
